#include "primer/trie.h"
#include <cstddef>
#include <iostream>
#include <memory>
#include <string_view>
#include <vector>
#include "common/exception.h"

namespace bustub {

template <class T>
auto Trie::Get(std::string_view key) const -> const T * {
  // throw NotImplementedException("Trie::Get is not implemented.");

  // You should walk through the trie to find the node corresponding to the key. If the node doesn't exist, return
  // nullptr. After you find the node, you should use `dynamic_cast` to cast it to `const TrieNodeWithValue<T> *`. If
  // dynamic_cast returns `nullptr`, it means the type of the value is mismatched, and you should return nullptr.
  // Otherwise, return the value.
  auto parent = root_;
  for (size_t i = 0; i < key.length() && parent; i++) {
    if (parent == nullptr) {
      return nullptr;
    }
    if (parent->children_.find(key[i]) == parent->children_.end()) {
      return nullptr;
    }
    if (i == key.length() - 1) {
      auto ret = parent->children_.at(key[i]);
      auto *cast = dynamic_cast<const TrieNodeWithValue<T> *>(ret.get());
      if (cast == nullptr) {
        return nullptr;
      }
      return cast->value_.get();
    }
    parent = parent->children_.at(key[i]);
  }
  auto *cast = dynamic_cast<const TrieNodeWithValue<T> *>(root_.get());
  if (cast == nullptr) {
    return nullptr;
  }
  return cast->value_.get();
}

template <class T>
auto Trie::Put(std::string_view key, T value) const -> Trie {
  // Note that `T` might be a non-copyable type. Always use `std::move` when creating `shared_ptr` on that value.
  // throw NotImplementedException("Trie::Put is not implemented.");

  // You should walk through the trie and create new nodes if necessary. If the node corresponding to the key already
  // exists, you should create a new `TrieNodeWithValue`.
  auto root = root_;
  if (root == nullptr) {
    root = std::make_shared<TrieNode>();
  }

  if (key.empty()) {
    root = std::make_shared<TrieNodeWithValue<T>>(
        TrieNodeWithValue<T>(root->children_, std::make_shared<T>(std::move(value))));
    return Trie(root);
  }

  auto parent = std::shared_ptr<TrieNode>(root->Clone());  // parent
  root = parent;
  for (size_t i = 0; i < key.length(); i++) {
    std::shared_ptr<TrieNode> new_node;
    if (parent == nullptr) {
      new_node = std::make_shared<TrieNode>();
      parent->children_[key[i]] = new_node;
    }
    if (i == key.length() - 1) {
      if (parent->children_.find(key[i]) != parent->children_.end()) {
        auto new_node_tmp = std::make_shared<TrieNodeWithValue<T>>(
            TrieNodeWithValue<T>(parent->children_[key[i]]->children_, std::make_shared<T>(std::move(value))));
        parent->children_[key[i]] = new_node_tmp;
        new_node = new_node_tmp;
      } else {
        auto new_node_tmp =
            std::make_shared<TrieNodeWithValue<T>>(TrieNodeWithValue<T>(std::make_shared<T>(std::move(value))));
        parent->children_[key[i]] = new_node_tmp;
        new_node = new_node_tmp;
      }
      return Trie(root);
    }
    if (parent->children_.find(key[i]) != parent->children_.end()) {
      new_node = std::shared_ptr<TrieNode>(parent->children_[key[i]]->Clone());
      parent->children_[key[i]] = new_node;
    } else {
      new_node = std::make_shared<TrieNode>();
      parent->children_[key[i]] = new_node;
    }
    parent = new_node;
  }
  return Trie(root);
}

auto Trie::Remove(std::string_view key) const -> Trie {
  // throw NotImplementedException("Trie::Remove is not implemented.");

  // You should walk through the trie and remove nodes if necessary. If the node doesn't contain a value any more,
  // you should convert it to `TrieNode`. If a node doesn't have children any more, you should remove it.
  if (root_ == nullptr) {
    return Trie(root_);
  }
  if (key.empty()) {
    return Trie(std::make_shared<const TrieNode>(root_->children_));
  }
  auto parent = std::shared_ptr<TrieNode>(root_->Clone());
  auto new_root = parent;

  // Copy
  for (char i : key) {
    if (parent == nullptr) {
      return Trie(root_);
    }
    if (parent->children_.find(i) == parent->children_.end()) {
      return Trie(root_);
    }
    parent = parent->children_[i]->Clone();
  }
  TryRemove(key, 0, new_root.get(), new_root.get());
  return Trie(new_root);
}

// 若尾结点是值结点，判断有没孩子，没孩子删除，递归往上，
// 碰到值结点返回，非值结点，判断有没孩子，没孩子删除，有孩子返回
// 有孩子设置为非值结点，返回
// 尾结点非值结点，返回
// 先实现原地删除，再实现COW
auto Trie::TryRemove(std::string_view key, size_t index, TrieNode *parent, TrieNode *old_parent) const -> void {
  parent->children_[key[index]] = old_parent->children_[key[index]]->Clone();
  if (index == key.size() - 1) {
    if (parent->children_[key[index]]->is_value_node_) {
      // 若尾结点是值结点，判断有没孩子，没孩子删除
      if (parent->children_[key[index]]->children_.empty()) {
        parent->children_.erase(key[index]);
      } else {  // 有孩子转化为非值结点，返回
        parent->children_[key[index]] = std::make_shared<TrieNode>(parent->children_[key[index]]->children_);
      }
    } else {
      return;
    }
    return;
  }
  const TrieNode *cur = parent->children_[key[index]].get();
  const TrieNode *old_cur = old_parent->children_[key[index]].get();
  TryRemove(key, index + 1, const_cast<TrieNode *>(cur), const_cast<TrieNode *>(old_cur));
  // 非尾结点碰到值结点返回
  if (cur->is_value_node_) {
    return;
  }
  // 非值非尾结点，判断有没孩子，有孩子返回，没孩子删除
  if (!cur->children_.empty()) {
    return;
  }
  parent->children_.erase(key[index]);
}

// Below are explicit instantiation of template functions.
//
// Generally people would write the implementation of template classes and functions in the header file. However, we
// separate the implementation into a .cpp file to make things clearer. In order to make the compiler know the
// implementation of the template functions, we need to explicitly instantiate them here, so that they can be picked up
// by the linker.

template auto Trie::Put(std::string_view key, uint32_t value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const uint32_t *;

template auto Trie::Put(std::string_view key, uint64_t value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const uint64_t *;

template auto Trie::Put(std::string_view key, std::string value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const std::string *;

// If your solution cannot compile for non-copy tests, you can remove the below lines to get partial score.

using Integer = std::unique_ptr<uint32_t>;

template auto Trie::Put(std::string_view key, Integer value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const Integer *;

template auto Trie::Put(std::string_view key, MoveBlocked value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const MoveBlocked *;

}  // namespace bustub
